package prom

import (
	"context"
	dto "github.com/prometheus/client_model/go"
	"github.com/prometheus/prom2json"
	"gluenet/pkg/log/v1"
	"net/http"
	"sync"
	"time"
)

type Collector struct {
	ctx      context.Context
	cancel   context.CancelFunc
	interval time.Duration
	url      string

	//mfl map[string]func(*dto.MetricFamily) float64
	mfl *sync.Map
	//rch map[string]chan float64
	rch *sync.Map
}

func NewCollector(ctx context.Context, interval time.Duration, url string) *Collector {
	ctx2, cancel := context.WithCancel(ctx)
	return &Collector{
		ctx:      ctx2,
		cancel:   cancel,
		interval: interval,
		url:      url,

		mfl: &sync.Map{},
		//mfl: make(map[string]func(*dto.MetricFamily) float64),
		rch: &sync.Map{},
		//rch: make(map[string]chan float64),
	}
}

func (r *Collector) AddLabel(fname string, f func(*dto.MetricFamily) float64) (chan float64, error) {
	if _, ok := r.mfl.Load(fname); ok {
		ch, _ := r.rch.Load(fname)
		return ch.(chan float64), nil
	}
	//r.mfl[fname] = f
	var ch = make(chan float64)
	r.rch.Store(fname, ch)
	r.mfl.Store(fname, f)

	return ch, nil
}

func (r *Collector) Run(logger v1.Logger) {
	tk := time.NewTicker(r.interval)
	for {
		select {
		case <-r.ctx.Done():
			return
		case <-tk.C:
			rec := make(chan *dto.MetricFamily, 1024)
			if err := prom2json.FetchMetricFamilies(r.url, rec, &http.Transport{}); err != nil {
				logger.Errorf("fetch error %v", err)
				continue
			}

			for mf := range rec {
				//fmt.Printf("%v\n", *mf.Name)
				r.mfl.Range(func(l, f interface{}) bool {
					if *mf.Name == l.(string) {
						chi, _ := r.rch.Load(l)
						ch := chi.(chan float64)
						ch <- f.(func(*dto.MetricFamily) float64)(mf)
					}
					return true
				})
			}
		}
	}
}

func (r *Collector) Close() {
	r.cancel()
}
